package info.wangyuan.agent.service.conversation;

import info.wangyuan.agent.utils.SseEmitterUtils;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 聊天会话管理器 负责管理正在进行的对话会话，支持会话中断功能
 */
@Slf4j
@Component
public class ChatSessionManager {

    /**
     * 会话信息
     */
    @Data
    public static class SessionInfo {
        private final String sessionId;
        private final SseEmitter emitter;
        private final AtomicBoolean interrupted;
        private final long startTime;

        // 绑定的流式任务（CompletableFuture）
        private volatile CompletableFuture<?> streamTask;

        public SessionInfo(String sessionId, SseEmitter emitter) {
            this.sessionId = sessionId;
            this.emitter = emitter;
            this.interrupted = new AtomicBoolean(false);
            this.startTime = System.currentTimeMillis();
        }

        public boolean isInterrupted() {
            return interrupted.get();
        }

        public void setInterrupted() {
            interrupted.set(true);
        }
    }

    // 使用sessionId作为key，存储正在进行的对话会话
    private final ConcurrentHashMap<String, SessionInfo> activeSessions = new ConcurrentHashMap<>();

    /**
     * 注册一个新的对话会话
     */
    public void registerSession(String sessionId, SseEmitter emitter) {
        SessionInfo sessionInfo = new SessionInfo(sessionId, emitter);
        activeSessions.put(sessionId, sessionInfo);
        log.info("注册对话会话: sessionId={}", sessionId);

        emitter.onCompletion(() -> {
            removeSession(sessionId);
            log.info("对话会话完成: sessionId={}", sessionId);
        });

        emitter.onTimeout(() -> {
            removeSession(sessionId);
            log.warn("对话会话超时: sessionId={}", sessionId);
        });

        emitter.onError((throwable) -> {
            removeSession(sessionId);
            log.error("对话会话错误: sessionId={}, error={}", sessionId, throwable.getMessage());
        });
    }

    /**
     * 绑定流式任务到会话
     */
    public void bindStream(String sessionId, CompletableFuture<?> streamTask) {
        SessionInfo sessionInfo = activeSessions.get(sessionId);
        if (sessionInfo != null) {
            sessionInfo.setStreamTask(streamTask);
            log.info("绑定流式任务到会话: sessionId={}", sessionId);
        } else {
            log.warn("尝试绑定流式任务到不存在的会话: sessionId={}", sessionId);
        }
    }

    /**
     * 移除对话会话
     */
    public void removeSession(String sessionId) {
        SessionInfo removed = activeSessions.remove(sessionId);
        if (removed != null) {
            long duration = System.currentTimeMillis() - removed.getStartTime();
            log.info("移除对话会话: sessionId={}, 持续时间={}ms", sessionId, duration);
        }
    }

    /**
     * 中断指定的对话会话
     */
    public boolean interruptSession(String sessionId) {
        SessionInfo sessionInfo = activeSessions.get(sessionId);
        if (sessionInfo == null) {
            log.warn("尝试中断不存在的会话: sessionId={}", sessionId);
            return false;
        }

        sessionInfo.setInterrupted();
        log.info("设置会话中断标志: sessionId={}", sessionId);

        activeSessions.remove(sessionId);

        try {
            // 取消流式任务
            CompletableFuture<?> task = sessionInfo.getStreamTask();
            if (task != null && !task.isDone()) {
                boolean cancelled = task.cancel(true);
                log.info("尝试取消流式任务: sessionId={}, 结果={}", sessionId, cancelled);
            }

            // 通知前端
            SseEmitter emitter = sessionInfo.getEmitter();
            SseEmitterUtils.safeSend(emitter,
                    SseEmitter.event().name("interrupt").data("{\"interrupted\": true, \"message\": \"对话已被中断\"}"));
            SseEmitterUtils.safeComplete(emitter);

            log.info("对话会话已中断: sessionId={}", sessionId);
            return true;

        } catch (Exception e) {
            log.error("中断会话时发生错误: sessionId={}, error={}", sessionId, e.getMessage());
            return true;
        }
    }

    public boolean isSessionInterrupted(String sessionId) {
        SessionInfo sessionInfo = activeSessions.get(sessionId);
        return sessionInfo != null && sessionInfo.isInterrupted();
    }

    public int getActiveSessionCount() {
        return activeSessions.size();
    }

    public boolean hasSession(String sessionId) {
        return activeSessions.containsKey(sessionId);
    }
}
